//
// Created by yanhai on 2020/7/15.
//

#include <iostream>
#include <fstream>
#include <chrono>
#include <vector>
#include <hs.h>
#include <climits>
#include <cstring>
#include <bitset>

using namespace std;

// Simple timing class
class Clock {
public:
    void start()
    {
        time_start = std::chrono::system_clock::now();
    }

    void stop()
    {
        time_end = std::chrono::system_clock::now();
    }

    double seconds() const
    {
        std::chrono::duration<double> delta = time_end - time_start;
        return delta.count();
    }

private:
    std::chrono::time_point<std::chrono::system_clock> time_start, time_end;
};

static void parseFile(const char *filename, vector<string> &patterns,
                      vector<unsigned> &flags, vector<unsigned> &ids);

static hs_database_t *buildDatabase(const vector<const char *> &expressions,
                                    const vector<unsigned> &flags,
                                    const vector<unsigned> &ids)
{
    hs_database_t *db;
    hs_compile_error_t *compileErr;
    hs_error_t err;

    Clock clock;
    clock.start();

    err = hs_compile_multi(expressions.data(), flags.data(), ids.data(),
                           expressions.size(), HS_MODE_BLOCK, nullptr, &db, &compileErr);

    clock.stop();

    if (err != HS_SUCCESS) {
        if (compileErr->expression < 0) {
            // The error does not refer to a particular expression.
            cerr << "ERROR: " << compileErr->message << endl;
        } else {
            cerr << "ERROR: Pattern '" << expressions[compileErr->expression]
                 << "' failed compilation with error: " << compileErr->message
                 << endl;
        }
        // As the compileErr pointer points to dynamically allocated memory, if
        // we get an error, we must be sure to release it. This is not
        // necessary when no error is detected.
        hs_free_compile_error(compileErr);
        exit(-1);
    }

    cout << "Hyperscan database compiled in "
         << clock.seconds() << " seconds."
         << endl;

    return db;
}

/**
 * This function will read in the file with the specified name, with an
 * expression per line, ignoring lines starting with '#' and build a Hyperscan
 * database for it.
 */
static hs_database_t *databasesFromFile(const char *filename, vector<string> &patterns,
                                        vector<unsigned> &flags, vector<unsigned> &ids)
{
    // do the actual file reading and string handling
    parseFile(filename, patterns, flags, ids);

    vector<const char *> cstrPatterns;
    for (const auto &pattern : patterns) {
        cstrPatterns.push_back(pattern.c_str());
    }

    return buildDatabase(cstrPatterns, flags, ids);
}

/**
 * Fill a data buffer from the given filename, returning it and filling @a
 * length with its length. Returns NULL on failure.
 */
static char *readInputData(const char *inputFN, unsigned int *length)
{
    FILE *f = fopen(inputFN, "rb");
    if (!f) {
        fprintf(stderr, "ERROR: unable to open file \"%s\": %s\n", inputFN,
                strerror(errno));
        return NULL;
    }

    /* We use fseek/ftell to get our data length, in order to keep this example
     * code as portable as possible. */
    if (fseek(f, 0, SEEK_END) != 0) {
        fprintf(stderr, "ERROR: unable to seek file \"%s\": %s\n", inputFN,
                strerror(errno));
        fclose(f);
        return NULL;
    }
    long dataLen = ftell(f);
    if (dataLen < 0) {
        fprintf(stderr, "ERROR: ftell() failed: %s\n", strerror(errno));
        fclose(f);
        return NULL;
    }
    if (fseek(f, 0, SEEK_SET) != 0) {
        fprintf(stderr, "ERROR: unable to seek file \"%s\": %s\n", inputFN,
                strerror(errno));
        fclose(f);
        return NULL;
    }

    /* Hyperscan's hs_scan function accepts length as an unsigned int, so we
     * limit the size of our buffer appropriately. */
    if ((unsigned long) dataLen > UINT_MAX) {
        dataLen = UINT_MAX;
        printf("WARNING: clipping data to %ld bytes\n", dataLen);
    } else if (dataLen == 0) {
        fprintf(stderr, "ERROR: input file \"%s\" is empty\n", inputFN);
        fclose(f);
        return NULL;
    }

    char *inputData = (char *) malloc(dataLen);
    if (!inputData) {
        fprintf(stderr, "ERROR: unable to malloc %ld bytes\n", dataLen);
        fclose(f);
        return NULL;
    }

    char *p = inputData;
    size_t bytesLeft = dataLen;
    while (bytesLeft) {
        size_t bytesRead = fread(p, 1, bytesLeft, f);
        bytesLeft -= bytesRead;
        p += bytesRead;
        if (ferror(f) != 0) {
            fprintf(stderr, "ERROR: fread() failed\n");
            free(inputData);
            fclose(f);
            return NULL;
        }
    }

    fclose(f);

    *length = (unsigned int) dataLen;
    return inputData;
}

struct CallbackCtx {
    const char *data;
    vector<string> &patterns;
    vector<unsigned> &flags;
    vector<unsigned> &ids;
};

/**
 * utf8编码规则
 * 1字节  0xxx xxxx
 * 2字节  110x xxxx | 10xx xxxx
 * 3字节  1110 xxxx | 10xx xxxx | 10xx xxxx
 * 4字节  1111 0xxx | 10xx xxxx | 10xx xxxx | 10xx xxxx
 * ...
 * @param byte 输入字节
 * @return 当字节是首字节时，返回字节长度，当字节是中间字节时，返回0
 */
static inline int Utf8ByteLen(char byte)
{
    bitset<8> bs(byte);
    // 0xxxxxxx 1字节长度的
    if (!bs.test(7)) {
        return 1;
    }
    // 10xxxxxx 多字节长度的中间字节
    if (bs.test(7) && !bs.test(6)) {
        return 0;
    }
    // 判断多字节长度
    for (int i = 7; i > 0; i--) {
        if (!bs.test(i)) {
            return 7 - i;
        }
    }

    return 0;
}

/**
 * 在utf8编码的文本中，截取一段
 * @param data 文本
 * @param offset 从offset开始截取
 * @param len 截取len长度
 * @return 截取后的文本，string类型
 */
static string splitUtf8(const char *data, unsigned long long offset, unsigned long long len)
{
    // 判断开始字节
    while (!Utf8ByteLen(data[offset])) {
        offset++;
    }
    // 判断结束字节
    while (!(Utf8ByteLen(data[offset + len]))) {
        len--;
    }

    return string(data + offset, len);
}

/**
 * 在回调函数中，打印关键字等的上下文，长度为128字节
 * @param id
 * @param from
 * @param to
 * @param flags
 * @param ctx
 * @return
 */
static int eventHandler(unsigned int id, unsigned long long from,
                        unsigned long long to, unsigned int flags, void *ctx)
{
    auto cctx = reinterpret_cast<CallbackCtx *>(ctx);
    cout << "Match for id \"" << id << "\", pattern: \"" << cctx->patterns[id].c_str()
         << "\", at from " << from << " offset " << to << ", flag: " << flags << endl;
    unsigned long long exp_len = to - from;
    unsigned int offset = 0;
    if (exp_len < 128) {
        offset = (128 - exp_len) / 2;
        exp_len = 128;
    }

    if (from > 0) {
        offset = from - offset;
    } else {
        offset = to - 64;
        exp_len = 128;
    }

    auto result = splitUtf8(cctx->data, offset, exp_len);
    cout << result.length() << ": " << result.c_str() << endl
         << "==============================" << endl << endl;

    return 0;
}

int main(int argc, char *argv[])
{
    if (argc != 3) {
        cerr << "Usage: " << argv[0] << " <pattern> <input file>" << endl;
        return -1;
    }

    const char *patternFile = argv[1];
    const char *textFile = argv[2];

    cout << "Pattern file: " << patternFile << endl;

    vector<string> patterns;
    vector<unsigned> flags;
    vector<unsigned> ids;

    hs_database_t *database = databasesFromFile(patternFile, patterns, flags, ids);

    /* Next, we read the input data file into a buffer. */
    unsigned int length;
    char *inputData = readInputData(textFile, &length);
    if (!inputData) {
        hs_free_database(database);
        return -1;
    }

    CallbackCtx cctx = {inputData, patterns, flags, ids};

    printf("Scanning %u bytes with Hyperscan\n", length);

    hs_scratch_t *scratch = nullptr;
    if (hs_alloc_scratch(database, &scratch) != HS_SUCCESS) {
        fprintf(stderr, "ERROR: Unable to allocate scratch space. Exiting.\n");
        free(inputData);
        hs_free_database(database);
        return -1;
    }

    if (hs_scan(database, inputData, length, 0, scratch, eventHandler,
                &cctx) != HS_SUCCESS) {
        fprintf(stderr, "ERROR: Unable to scan input buffer. Exiting.\n");
        hs_free_scratch(scratch);
        free(inputData);
        hs_free_database(database);
        return -1;
    }

    /* Scanning is complete, any matches have been handled, so now we just
     * clean up and exit.
     */
    hs_free_scratch(scratch);
    free(inputData);
    hs_free_database(database);
    return 0;
}

static unsigned parseFlags(const string &flagsStr)
{
    unsigned flags = 0;
    for (const auto &c : flagsStr) {
        switch (c) {
            case 'i':
                flags |= HS_FLAG_CASELESS;
                break;
            case 'm':
                flags |= HS_FLAG_MULTILINE;
                break;
            case 's':
                flags |= HS_FLAG_DOTALL;
                break;
            case 'H':
                flags |= HS_FLAG_SINGLEMATCH;
                break;
            case 'V':
                flags |= HS_FLAG_ALLOWEMPTY;
                break;
            case '8':
                flags |= HS_FLAG_UTF8;
                break;
            case 'W':
                flags |= HS_FLAG_UCP;
                break;
            case 'C':
                flags |= HS_FLAG_COMBINATION;
                break;
            case 'Q':
                flags |= HS_FLAG_QUIET;
                break;
            case 'L':
                flags |= HS_FLAG_SOM_LEFTMOST;
                break;
            case '\r': // stray carriage-return
                break;
            default:
                cerr << "Unsupported flag \'" << c << "\'" << endl;
                exit(-1);
        }
    }
    return flags;
}

static void parseFile(const char *filename, vector<string> &patterns,
                      vector<unsigned> &flags, vector<unsigned> &ids)
{
    ifstream inFile(filename);
    if (!inFile.good()) {
        cerr << "ERROR: Can't open pattern file \"" << filename << "\"" << endl;
        exit(-1);
    }

    for (unsigned i = 1; !inFile.eof(); ++i) {
        string line;
        getline(inFile, line);

        // if line is empty, or a comment, we can skip it
        if (line.empty() || line[0] == '#') {
            continue;
        }

        // otherwise, it should be ID:PCRE, e.g.
        //  10001:/foobar/is

        size_t colonIdx = line.find_first_of(':');
        if (colonIdx == string::npos) {
            cerr << "ERROR: Could not parse line " << i << endl;
            exit(-1);
        }

        // we should have an unsigned int as an ID, before the colon
        unsigned id = std::stoi(line.substr(0, colonIdx));

        // rest of the expression is the PCRE
        const string expr(line.substr(colonIdx + 1));

        size_t flagsStart = expr.find_last_of('/');
        if (flagsStart == string::npos) {
            cerr << "ERROR: no trailing '/' char" << endl;
            exit(-1);
        }

        string pcre(expr.substr(1, flagsStart - 1));
        string flagsStr(expr.substr(flagsStart + 1, expr.size() - flagsStart));
        unsigned flag = parseFlags(flagsStr);

        patterns.push_back(pcre);
        flags.push_back(flag);
        ids.push_back(id);
    }
}
